package com.ververica.cdc.guass.sink.jdbc.internal.executor;

import com.ververica.cdc.guass.Constants;
import com.ververica.cdc.guass.source.kafka.data.NullableValue;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.table.data.*;
import org.apache.flink.table.data.binary.BinaryRowData;
import org.apache.flink.table.types.logical.DecimalType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.TimestampType;
import org.apache.flink.types.RowKind;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.SQLException;
import java.time.Instant;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

/**
 * TableBufferReducedStatementExecutor 的作用是缓存和优化 INSERT、UPDATE 和 DELETE 事件。
 * 它内部维护了一个缓存结构 reduceBuffer，用于存储待处理的事件。当调用 addToBatch 方法添加事件时，
 * 会将事件按照主键进行缓存和合并。当调用 executeBatch 方法执行批处理时，
 * 会根据缓存中的事件类型（INSERT/UPDATE 或 DELETE），分别调用相应的语句执行器进行批量执行。
 */
public final class TableBufferReducedStatementExecutor implements JdbcBatchStatementExecutor<RowData> {

    private final JdbcBatchStatementExecutor<RowData> upsertExecutor;
    private final JdbcBatchStatementExecutor<RowData> deleteExecutor;
    private final Function<RowData, RowData> keyExtractor;
    private static final int BATCH_SIZE = 2; // 定义批次大小

    private static final int BATCH_SIZE_THRESHOLD = 2; // 可以根据需求调整

    private final LogicalType[] fieldTypes;
    private final String[] fieldNames;

    // the mapping is [KEY, <+/-, VALUE>]
    private final Map<RowData, Tuple2<Boolean, RowData>> reduceBuffer = new HashMap<>();

    private static final Logger LOG = LoggerFactory.getLogger(TableBufferReducedStatementExecutor.class);


    public TableBufferReducedStatementExecutor(JdbcBatchStatementExecutor<RowData> upsertExecutor, JdbcBatchStatementExecutor<RowData> deleteExecutor, Function<RowData, RowData> keyExtractor, LogicalType[] fieldTypes, String[] fieldNames) {
        this.upsertExecutor = upsertExecutor;
        this.deleteExecutor = deleteExecutor;
        this.keyExtractor = keyExtractor;
        this.fieldTypes = fieldTypes;
        this.fieldNames = fieldNames;
    }

    @Override
    public void prepareStatements(Connection connection) throws SQLException {
        upsertExecutor.prepareStatements(connection);
        deleteExecutor.prepareStatements(connection);
    }

    @Override
    public void addToBatch(RowData record) throws SQLException {

        // Convert BinaryRowData to GenericRowData
        GenericRowData genericRowData = convertToGenericRowData(record, fieldTypes, fieldNames);

        RowData key = keyExtractor.apply(genericRowData);
        boolean flag = changeFlag(genericRowData.getRowKind());
        reduceBuffer.put(key, Tuple2.of(flag, genericRowData));


/*        //打印resultGenericRowData
        StringBuffer sb = new StringBuffer();

        for (int i = 0; i < genericRowData.getArity(); i++) {
            sb.append(((NullableValue) genericRowData.getField(i)).getValue() + "");
        }
        sb = sb.deleteCharAt(sb.length() - 1);
        LOG.info("genericRowData is {}", sb);*/
    }


    /**
     * Converts a RowData instance (expected to be BinaryRowData) to GenericRowData based on fieldTypes.
     *
     * @param rowData    The input RowData instance.
     * @param fieldTypes The LogicalType array representing the schema.
     * @param fieldNames
     * @return A CustomRowData instance populated with the field values.
     * @throws SQLException If the rowData is not an instance of BinaryRowData or if an unsupported type is encountered.
     */

    private GenericRowData convertToGenericRowData(RowData rowData, LogicalType[] fieldTypes, String[] fieldNames) throws SQLException {
        GenericRowData resultGenericRowData;
        String marks = null;
        Integer marksIndex = -1;

        // 查找 marks 字段
        for (int i = 0; i < fieldNames.length; i++) {
            if (fieldNames[i].toUpperCase().contains(Constants.PX_GAUSS_MARKS.toUpperCase())) {
                marksIndex = i;
                break;
            }
        }

        if (marksIndex == -1) {
            throw new RuntimeException("No px_gauss_marks found");
        }

        if (rowData instanceof GenericRowData) {
            GenericRowData inputGenericRowData = (GenericRowData) rowData;
            resultGenericRowData = new GenericRowData(inputGenericRowData.getRowKind(), fieldTypes.length);

            // 从输入的 GenericRowData 中获取 marks
            Object marksField = inputGenericRowData.getField(marksIndex);
            if (marksField instanceof StringData) {
                marks = ((StringData) marksField).toString();
            }

            for (int i = 0; i < fieldTypes.length; i++) {
                if (i == marksIndex) continue;

                Object fieldValue = inputGenericRowData.getField(i);
                boolean exists = marks != null && marks.charAt(i) != '0';
                boolean isNull = marks != null && marks.charAt(i) == '2';

                if (exists && !isNull) {
                    // 转换字段值为适当的类型
                    fieldValue = inputGenericRowData.getField(i);
                } else {
                    fieldValue = null;
                }

                NullableValue nullableValue = new NullableValue(fieldValue, exists);
                resultGenericRowData.setField(i, nullableValue);
            }
        } else if (rowData instanceof BinaryRowData) {
            BinaryRowData binaryRowData = (BinaryRowData) rowData;
            resultGenericRowData = new GenericRowData(binaryRowData.getRowKind(), fieldTypes.length);

            marks = binaryRowData.getString(marksIndex).toString();

            for (int i = 0; i < fieldTypes.length; i++) {
                if (i == marksIndex) continue;

                LogicalType fieldType = fieldTypes[i];
                Object fieldValue = null;
                boolean exists = marks.charAt(i) != '0';
                boolean isNull = marks.charAt(i) == '2';

                if (exists && !isNull) {
                    fieldValue = convertBinaryField(binaryRowData, i, fieldType, fieldNames);
                }

                NullableValue nullableValue = new NullableValue(fieldValue, exists);
                resultGenericRowData.setField(i, nullableValue);
            }
        } else {
            throw new SQLException("Unsupported RowData type: " + rowData.getClass().getName());
        }

        // 设置 marks 字段
        resultGenericRowData.setField(marksIndex, new NullableValue(null, false));

        return resultGenericRowData;
    }


    private Object convertBinaryField(BinaryRowData binaryRowData, int index, LogicalType fieldType, String[] fieldNames) throws SQLException {
        switch (fieldType.getTypeRoot()) {
            case BOOLEAN:
                return binaryRowData.getBoolean(index);
            case TINYINT:
                return binaryRowData.getByte(index);
            case SMALLINT:
                return binaryRowData.getShort(index);
            case INTEGER:
                return binaryRowData.getInt(index);
            case BIGINT:
                return binaryRowData.getLong(index);
            case FLOAT:
                return binaryRowData.getFloat(index);
            case DOUBLE:
                return binaryRowData.getDouble(index);
            case CHAR:
            case VARCHAR:
                return StringData.fromString(binaryRowData.getString(index).toString());
            case DECIMAL:
                if (!(fieldType instanceof DecimalType)) {
                    throw new SQLException("Expected DecimalType for DECIMAL LogicalType");
                }
                DecimalType decimalType = (DecimalType) fieldType;
                return binaryRowData.getDecimal(index, decimalType.getPrecision(), decimalType.getScale());

            case DATE:
                Integer epochDay = binaryRowData.getInt(index);
                // 将 epochDay 转换为 LocalDate
                java.time.LocalDate convertedLocalDate = java.time.LocalDate.ofEpochDay(epochDay);
                // 将 LocalDate 转换为 java.sql.Date
                java.sql.Date sqlDate = java.sql.Date.valueOf(convertedLocalDate);
                return sqlDate;
            case TIMESTAMP_WITHOUT_TIME_ZONE:
                if (!(fieldType instanceof TimestampType)) {
                    throw new SQLException("Expected TimestampType for TIMESTAMP_WITHOUT_TIME_ZONE LogicalType");
                }
                TimestampType timestampType = (TimestampType) fieldType;
                TimestampData timestampWithoutTZ = binaryRowData.getTimestamp(index, timestampType.getPrecision());

                LOG.info("TIMESTAMP_WITHOUT_TIME_ZONE value: {}", timestampWithoutTZ);

                return binaryRowData.getTimestamp(index, timestampType.getPrecision());
            case TIMESTAMP_WITH_TIME_ZONE:
                if (!(fieldType instanceof TimestampType)) {
                    throw new SQLException("Expected TimestampType for TIMESTAMP_WITH_TIME_ZONE LogicalType");
                }
                TimestampType timestampWithTZType = (TimestampType) fieldType;
                long timestampMillis = binaryRowData.getTimestamp(index, timestampWithTZType.getPrecision()).getMillisecond();
                ZoneOffset zoneOffset = ZoneOffset.UTC;
                OffsetDateTime offsetDateTime = Instant.ofEpochMilli(timestampMillis).atOffset(zoneOffset);
                LOG.info("TIMESTAMP_WITH_TIME_ZONE value: {}", offsetDateTime);
                return offsetDateTime;
            case BINARY:
                return binaryRowData.getBinary(index);
            case ARRAY:
                return binaryRowData.getArray(index);
            case MAP:
                return binaryRowData.getMap(index);
            case ROW:
                if (!(fieldType instanceof RowType)) {
                    throw new SQLException("Expected RowType for ROW LogicalType");
                }
                RowType nestedRowType = (RowType) fieldType;
                return convertToGenericRowData(binaryRowData.getRow(index, nestedRowType.getFieldCount()), nestedRowType.getChildren().toArray(new LogicalType[0]), fieldNames);
            default:
                throw new SQLException("Unsupported LogicalType: " + fieldType.getTypeRoot());
        }
    }


    /**
     * Returns true if the row kind is INSERT or UPDATE_AFTER, returns false if the row kind is
     * DELETE or UPDATE_BEFORE.
     */
    private boolean changeFlag(RowKind rowKind) {
        switch (rowKind) {
            case INSERT:
            case UPDATE_AFTER:
                return true;
            case DELETE:
            case UPDATE_BEFORE:
                return false;
            default:
                throw new UnsupportedOperationException(String.format("Unknown row kind, the supported row kinds is: INSERT, UPDATE_BEFORE, UPDATE_AFTER," + " DELETE, but get: %s.", rowKind));
        }
    }

    @Override
    public void executeBatch() throws SQLException {

        //打印entry
        LOG.info("reduceBuffer: {}", reduceBuffer);

        if (!reduceBuffer.isEmpty()) {
            for (Map.Entry<RowData, Tuple2<Boolean, RowData>> entry : reduceBuffer.entrySet()) {

                if (entry.getValue().f0) {
                    upsertExecutor.addToBatch(entry.getValue().f1);
                } else {
                    // delete by key
                    deleteExecutor.addToBatch(entry.getKey());
                }
            }
            upsertExecutor.executeBatch();
            deleteExecutor.executeBatch();
            reduceBuffer.clear();
        }
    }

    @Override
    public void closeStatements() throws SQLException {
        upsertExecutor.closeStatements();
        deleteExecutor.closeStatements();
    }
}
